//    Copyright (C) 2012-2022 OpenVPN Inc.

// Base class for TCP link objects.

#ifndef OPENVPN_TRANSPORT_COMMONLINK_H
#define OPENVPN_TRANSPORT_COMMONLINK_H

#include <deque>
#include <utility> // for std::move
#include <memory>

#include <openvpn/io/io.hpp>

#include <openvpn/common/size.hpp>
#include <openvpn/common/rc.hpp>
#include <openvpn/common/socktypes.hpp>
#include <openvpn/error/excode.hpp>
#include <openvpn/frame/frame.hpp>
#include <openvpn/log/sessionstats.hpp>
#include <openvpn/transport/tcplinkbase.hpp>
#include <openvpn/transport/pktstream.hpp>
#include <openvpn/transport/mutate.hpp>

#ifdef OPENVPN_GREMLIN
#include <openvpn/transport/gremlin.hpp>
#endif

#if defined(OPENVPN_DEBUG_TCPLINK) && OPENVPN_DEBUG_TCPLINK >= 1
#define OPENVPN_LOG_TCPLINK_ERROR(x) OPENVPN_LOG(x)
#else
#define OPENVPN_LOG_TCPLINK_ERROR(x)
#endif

#if defined(OPENVPN_DEBUG_TCPLINK) && OPENVPN_DEBUG_TCPLINK >= 3
#define OPENVPN_LOG_TCPLINK_VERBOSE(x) OPENVPN_LOG(x)
#else
#define OPENVPN_LOG_TCPLINK_VERBOSE(x)
#endif

namespace openvpn {
namespace TCPTransport {

template <typename Protocol,
          typename ReadHandler,
          bool RAW_MODE_ONLY>
class LinkCommon : public LinkBase
{
    typedef std::deque<BufferPtr> Queue;

  public:
    typedef RCPtr<LinkCommon<Protocol, ReadHandler, RAW_MODE_ONLY>> Ptr;
    typedef Protocol protocol;

    // In raw mode, data is sent and received without any special encapsulation.
    // In non-raw mode, data is packetized by prepending a 16-bit length word
    // onto each packet.  The OpenVPN and DNS protocols run in non-raw mode,
    // while other TCP protocols such as HTTP or HTTPS would run in raw mode.
    // This method is a no-op if RAW_MODE_ONLY is true.
    void set_raw_mode(const bool mode)
    {
        set_raw_mode_read(mode);
        set_raw_mode_write(mode);
    }

    void set_raw_mode_read(const bool mode)
    {
        if (RAW_MODE_ONLY)
            raw_mode_read = true;
        else
            raw_mode_read = mode;
    }

    void set_raw_mode_write(const bool mode)
    {
        if (RAW_MODE_ONLY)
            raw_mode_write = true;
        else
            raw_mode_write = mode;
    }

    void set_mutate(const TransportMutateStream::Ptr &mutate_arg)
    {
        mutate = mutate_arg;
    }

    bool send_queue_empty() const
    {
        return send_queue_size() == 0;
    }

    void inject(const Buffer &src)
    {
        const size_t size = src.size();
        OPENVPN_LOG_TCPLINK_VERBOSE("TCP inject size=" << size);
        if (size && !RAW_MODE_ONLY)
        {
            BufferAllocated buf;
            frame_context.prepare(buf);
            buf.write(src.c_data(), size);
            BufferAllocated pkt;
            put_pktstream(buf, pkt);
        }
    }

    void start()
    {
        if (!halt)
            queue_recv(nullptr);
    }

    void stop()
    {
        halt = true;
#ifdef OPENVPN_GREMLIN
        if (gremlin)
            gremlin->stop();
#endif
    }

    void reset_align_adjust(const size_t align_adjust)
    {
        frame_context.reset_align_adjust(align_adjust + (is_raw_mode() ? 0 : 2));
    }

    unsigned int send_queue_size() const
    {
        return queue.size()
#ifdef OPENVPN_GREMLIN
               + (gremlin ? gremlin->send_size() : 0)
#endif
            ;
    }

    bool send(BufferAllocated &b)
    {
        if (halt)
        {
            return false;
        }
        if (send_queue_max_size && send_queue_size() >= send_queue_max_size)
        {
            stats->error(Error::TCP_OVERFLOW);
            read_handler->tcp_error_handler("TCP_OVERFLOW");
            stop();
            return false;
        }

        BufferPtr buf;
        if (!free_list.empty())
        {
            buf = free_list.front();
            free_list.pop_front();
        }
        else
        {
            buf.reset(new BufferAllocated());
        }
        buf->swap(b);
        if (!is_raw_mode_write())
        {
            PacketStream::prepend_size(*buf);
        }
        if (mutate)
        {
            mutate->pre_send(*buf);
        }
#ifdef OPENVPN_GREMLIN
        if (gremlin)
        {
            gremlin_queue_send_buffer(buf);
        }
        else
#endif
        {
            from_app_send_buffer(buf);
        }
        return true;
    }

    void queue_recv(PacketFrom *tcpfrom)
    {
        OPENVPN_LOG_TCPLINK_VERBOSE("TLSLink::queue_recv");
        if (!tcpfrom)
            tcpfrom = new PacketFrom();
        frame_context.prepare(tcpfrom->buf);

        socket.async_receive(frame_context.mutable_buffer_clamp(tcpfrom->buf),
                             [self = Ptr(this), tcpfrom = PacketFrom::SPtr(tcpfrom)](const openvpn_io::error_code &error, const size_t bytes_recvd) mutable
                             {
            OPENVPN_ASYNC_HANDLER;
            try
            {
                self->handle_recv(std::move(tcpfrom), error, bytes_recvd);
            }
            catch (const std::exception &e)
            {
                Error::Type err = Error::TCP_SIZE_ERROR;
                const char *msg = "TCP_SIZE_ERROR";
                // if exception is an ExceptionCode, translate the code
                // to return status string
                {
                    const ExceptionCode *ec = dynamic_cast<const ExceptionCode *>(&e);
                    if (ec && ec->code_defined())
                    {
                        err = ec->code();
                        msg = ec->what();
                    }
                }

                OPENVPN_LOG_TCPLINK_ERROR("TCP packet extract exception: " << e.what());
                self->stats->error(err);
                self->read_handler->tcp_error_handler(msg);
                self->stop();
            }
        });
    }

  protected:
    LinkCommon(ReadHandler read_handler_arg,
               typename Protocol::socket &socket_arg,
               const size_t send_queue_max_size_arg, // 0 to disable
               const size_t free_list_max_size_arg,
               const Frame::Context &frame_context_arg,
               const SessionStats::Ptr &stats_arg)
        : socket(socket_arg),
          read_handler(read_handler_arg),
          frame_context(frame_context_arg),
          stats(stats_arg),
          send_queue_max_size(send_queue_max_size_arg),
          free_list_max_size(free_list_max_size_arg)
    {
        set_raw_mode(false);
    }

#ifdef OPENVPN_GREMLIN
    void gremlin_config(const Gremlin::Config::Ptr &config)
    {
        if (config)
            gremlin.reset(new Gremlin::SendRecvQueue(socket.get_executor().context(), config, true));
    }
#endif

    bool is_raw_mode() const
    {
        return is_raw_mode_read() && is_raw_mode_write();
    }

    bool is_raw_mode_read() const
    {
        if (RAW_MODE_ONLY)
            return true;
        else
            return raw_mode_read;
    }

    bool is_raw_mode_write() const
    {
        if (RAW_MODE_ONLY)
            return true;
        else
            return raw_mode_write;
    }

    LinkCommon()
    {
        stop();
    }

    void queue_send_buffer(BufferPtr &buf)
    {
        queue.push_back(std::move(buf));
        if (queue.size() == 1) // send operation not currently active?
            queue_send();
    }

    void queue_send()
    {
        BufferAllocated &buf = *queue.front();
        socket.async_send(buf.const_buffer_clamp(),
                          [self = Ptr(this)](const openvpn_io::error_code &error, const size_t bytes_sent)
                          {
            OPENVPN_ASYNC_HANDLER;
            self->handle_send(error, bytes_sent);
        });
    }

    void handle_send(const openvpn_io::error_code &error, const size_t bytes_sent)
    {
        if (!halt)
        {
            if (!error)
            {
                OPENVPN_LOG_TCPLINK_VERBOSE("TLS-TCP send raw=" << raw_mode_write << " size=" << bytes_sent);
                stats->inc_stat(SessionStats::BYTES_OUT, bytes_sent);
                stats->inc_stat(SessionStats::PACKETS_OUT, 1);

                BufferPtr buf = queue.front();
                if (bytes_sent == buf->size())
                {
                    queue.pop_front();
                    if (free_list.size() < free_list_max_size)
                    {
                        buf->reset_content();
                        free_list.push_back(std::move(buf)); // recycle the buffer for later use
                    }
                }
                else if (bytes_sent < buf->size())
                    buf->advance(bytes_sent);
                else
                {
                    stats->error(Error::TCP_OVERFLOW);
                    read_handler->tcp_error_handler("TCP_INTERNAL_ERROR"); // error sent more bytes than we asked for
                    stop();
                    return;
                }
            }
            else
            {
                OPENVPN_LOG_TCPLINK_ERROR("TLS-TCP send error: " << error.message());
                stats->error(Error::NETWORK_SEND_ERROR);
                read_handler->tcp_error_handler("NETWORK_SEND_ERROR");
                stop();
                return;
            }
            if (!queue.empty())
                queue_send();
            else
                tcp_write_queue_needs_send();
        }
    }

    bool process_recv_buffer(BufferAllocated &buf)
    {
        bool requeue = true;

        OPENVPN_LOG_TCPLINK_VERBOSE("TLSLink::process_recv_buffer: size=" << buf.size());

        if (!is_raw_mode_read())
        {
            try
            {
                BufferAllocated pkt;
                requeue = put_pktstream(buf, pkt);
                if (!buf.allocated() && pkt.allocated()) // recycle pkt allocated buffer
                    buf.move(pkt);
            }
            catch (const std::exception &e)
            {
                OPENVPN_LOG_TCPLINK_ERROR("TLS-TCP packet extract error: " << e.what());
                stats->error(Error::TCP_SIZE_ERROR);
                read_handler->tcp_error_handler("TCP_SIZE_ERROR");
                stop();
                return false;
            }
        }
        else
        {
            if (mutate)
            {
                mutate->post_recv(buf);
            }
#ifdef OPENVPN_GREMLIN
            if (gremlin)
            {
                requeue = gremlin_recv(buf);
            }
            else
#endif
            {
                requeue = read_handler->tcp_read_handler(buf);
            }
        }
        return requeue;
    }

    void handle_recv(PacketFrom::SPtr pfp, const openvpn_io::error_code &error, const size_t bytes_recvd)
    {
        OPENVPN_LOG_TCPLINK_VERBOSE("Link::handle_recv: " << error.message());
        if (!halt)
        {
            if (!error)
            {
                recv_buffer(pfp, bytes_recvd);
            }
            else if (error == openvpn_io::error::eof)
            {
                OPENVPN_LOG_TCPLINK_ERROR("TCP recv EOF");
                read_handler->tcp_eof_handler();
            }
            else
            {
                OPENVPN_LOG_TCPLINK_ERROR("TCP recv error: " << error.message());
                stats->error(Error::NETWORK_RECV_ERROR);
                read_handler->tcp_error_handler("NETWORK_RECV_ERROR");
                stop();
            }
        }
    }

    bool put_pktstream(BufferAllocated &buf, BufferAllocated &pkt)
    {
        bool requeue = true;
        stats->inc_stat(SessionStats::BYTES_IN, buf.size());
        stats->inc_stat(SessionStats::PACKETS_IN, 1);
        if (mutate)
            mutate->post_recv(buf);
        while (buf.size())
        {
            pktstream.put(buf, frame_context);
            if (pktstream.ready())
            {
                pktstream.get(pkt);
#ifdef OPENVPN_GREMLIN
                if (gremlin)
                    requeue = gremlin_recv(pkt);
                else
#endif
                    requeue = read_handler->tcp_read_handler(pkt);
            }
        }
        return requeue;
    }

#ifdef OPENVPN_GREMLIN
    void gremlin_queue_send_buffer(BufferPtr &buf)
    {
        gremlin->send_queue([self = Ptr(this), buf = std::move(buf)]() mutable
                            {
	    if (!self->halt)
	      {
		self->queue_send_buffer(buf);
	      } });
    }

    bool gremlin_recv(BufferAllocated &buf)
    {
        gremlin->recv_queue([self = Ptr(this), buf = std::move(buf)]() mutable
                            {
	    if (!self->halt)
	      {
		const bool requeue = self->read_handler->tcp_read_handler(buf);
		if (requeue)
		  self->queue_recv(nullptr);
	      } });
        return false;
    }
#endif

    void tcp_write_queue_needs_send()
    {
        read_handler->tcp_write_queue_needs_send();
    }

    typename Protocol::socket &socket;
    ReadHandler read_handler;
    Frame::Context frame_context;
    SessionStats::Ptr stats;
    const size_t send_queue_max_size;
    const size_t free_list_max_size;
    Queue queue;     // send queue
    Queue free_list; // recycled free buffers for send queue
    PacketStream pktstream;
    TransportMutateStream::Ptr mutate;
    bool raw_mode_read;
    bool raw_mode_write;
    bool halt = false;

#ifdef OPENVPN_GREMLIN
    std::unique_ptr<Gremlin::SendRecvQueue> gremlin;
#endif

  private:
    virtual void recv_buffer(PacketFrom::SPtr &pfp, const size_t bytes_recvd) = 0;
    virtual void from_app_send_buffer(BufferPtr &buf) = 0;
};
} // namespace TCPTransport
} // namespace openvpn

#endif
